#!/usr/bin/perl
#Programmer: Rory Carmichael
#Purpose: Create a random reverse translation of a given min-max input
#Usage: rc-rrt-gen.pl <input_fa_orfs> <codon_frequency_table>

if ($#ARGV != 1) {
	print STDERR "USAGE: rc-rrt-gen.pl <input_fa_orfs> <codon_frequency_table>\n";
	exit 1;
}

my %seqs = ();
my %codon_freqs = ();
my %codon_dice = ();
my $qry_file = $ARGV[0];
my $codon_file = $ARGV[1];

my %DNAtoAA = ('GCT' => 'A', 'GCC' => 'A', 'GCA' => 'A', 'GCG' => 'A', 'TGT' => 'C',
	       'TGC' => 'C', 'GAT' => 'D', 'GAC' => 'D', 'GAA' => 'E', 'GAG' => 'E',
	       'TTT' => 'F', 'TTC' => 'F', 'GGT' => 'G', 'GGC' => 'G', 'GGA' => 'G',
	       'GGG' => 'G', 'CAT' => 'H', 'CAC' => 'H', 'ATT' => 'I', 'ATC' => 'I',
	       'ATA' => 'I', 'AAA' => 'K', 'AAG' => 'K', 'TTG' => 'L', 'TTA' => 'L',
	       'CTT' => 'L', 'CTC' => 'L', 'CTA' => 'L', 'CTG' => 'L', 'ATG' => 'M',
	       'AAT' => 'N', 'AAC' => 'N', 'CCT' => 'P', 'CCC' => 'P', 'CCA' => 'P',
	       'CCG' => 'P', 'CAA' => 'Q', 'CAG' => 'Q', 'CGT' => 'R', 'CGC' => 'R',
	       'CGA' => 'R', 'CGG' => 'R', 'AGA' => 'R', 'AGG' => 'R', 'TCT' => 'S',
	       'TCC' => 'S', 'TCA' => 'S', 'TCG' => 'S', 'AGT' => 'S', 'AGC' => 'S',
	       'ACT' => 'T', 'ACC' => 'T', 'ACA' => 'T', 'ACG' => 'T', 'GTT' => 'V',
	       'GTC' => 'V', 'GTA' => 'V', 'GTG' => 'V', 'TGG' => 'W', 'TAT' => 'Y',
	       'TAC' => 'Y', 'TAA' => '*', 'TAG' => '*', 'TGA' => '*');

#HASH containing the number of synonymous codons for each amino acid, calculated because I am too lazy to enumerate them manually and that is why we invented computers
my %syncounts = ();
my %syns = ();
foreach my $key (keys(%DNAtoAA)) {
	$syncounts{$DNAtoAA{$key}}++;
	my @cursyn = @{$syns{$DNAtoAA{$key}}};
	push(@cursyn,$key);
	$syns{$DNAtoAA{$key}} = \@cursyn;
}	

sub read_codon_table {
	my $tfile = shift;
	print STDERR "Reading codon table from $tfile\n";
	open(IF, $tfile);
	foreach my $line (<IF>) {
		my @triplets = split(/\t+/,$line);
		foreach my $trip (@triplets) {
			my $codon = [split(/\s+/,$trip)]->[0];
			$freq = [ split(/\(/,$trip)]->[-1];
			$freq =~ s/\)//;
			$codon_freqs{$codon} = $freq;
		}
	}
	close(IF);
}

sub read_seqs {
	my $sfile = shift;
	print STDERR "Reading sequences from $sfile\n";
	open(IF, $sfile);
	my $curid;
	foreach my $line (<IF>) {
		chomp $line;
		if ($line =~ /^>/) {
			$curid = $line;
		} else {
			if (defined($curid)) {
				$line = "\U$line";
				$seqs{$curid} .= $line;
			} else {
				print STDERR "Sequence without id\n";
			}
		}
	}
	close(IF);
	while (my ($key,$val) = each(%seqs) ) {
		if (length($val) % 3  != 0 ) {
			print STDERR "Removing bad sequence $key of length " . length($val) . "\n";
			undef($seqs{$key});
		}
	}
}

sub make_codon_dice {
	while (my ($key,$val) = each(%syns)) {
		my %curdice = ();
		my $sum = 0;
		#sum up codons
		foreach my $codon (@$val) {
			$sum += $codon_freqs{$codon};
		}
#		print "$key sum is: $sum\n";
		if ($sum == 0) {
			print "STDERR: Amino Acid $key not present in query sequences\n";
			next;
		}
		#hash each synonymous codon on the percentage of time it appears
		my $lastval = 0;
		foreach my $codon (@$val) {
			my $key = $codon_freqs{$codon}/$sum;
			$lastval += $key;
			$curdice{$lastval} = $codon;
#			print "\tinitialized $codon with prob " . $codon_freqs{$codon}/$sum . " to $lastval \n";
		}
		#put the dice for this protein into the master hash
		$codon_dice{$key} = \%curdice;
	}
}

sub roll_codon_dice {
	my $AA = shift;
	my $roll = rand();
	foreach my $key (sort(keys(%{$codon_dice{$AA}}))) {
		if ($roll <= $key) {
			return $codon_dice{$AA}->{$key};
		}
	}
}

sub print_syns {
	while (my ($key,$val) = each(%syns)) {
		print "$key with $syncounts{$key} has codons:\t";
		foreach my $codon (@$val) {
			print "$codon ";
		}
		print "\n";
	}
}

sub reroll_seq {
	my $seqid = shift;
	my @seqarray = split(//,$seqs{$seqid});
	my $curcodon;
	print $seqid . "\n";
	foreach my $char (@seqarray) {
		$curcodon .= $char;
		if(length($curcodon) == 3) {
			print &roll_codon_dice($DNAtoAA{$curcodon});
			undef($curcodon);
		} 
	}
	print "\n";
}

#&print_syns();
&read_codon_table($codon_file);
&read_seqs($qry_file);
&make_codon_dice();
foreach my $key (keys(%seqs)) {
	&reroll_seq($key);
}
